from torch import nn
from transformers import AutoModel

from configuration import config

class product_classifier(nn.Module):
    #定义一个名为 product_classifier 的类，继承自 PyTorch 的 nn.Module 类，表示这是一个神经网络模型。
    def __init__(self,freeze_bert=True):
        """
        初始化模型，定义模型结构。
        参数：
            freeze_bert (bool): 是否冻结BERT的参数，默认为True。
        """
        super.__init__() #继承父类的初始化方法
        self.bert = AutoModel.from_pretrained(config.model_name) #加载预训练的BERT模型
        if freeze_bert:
            #冻结BERT的参数
            for param in self.bert.parameters(): #遍历BERT的参数
                param.requires_grad = False #设置参数的 requires_grad 属性为 False，表示不更新参数
        self.linear = nn.Linear(self.bert.config.hidden_size, config.NUM_CLASSES)
        #定义一个全连接层，输入维度为BERT的隐藏层维度，输出维度为 config.NUM_CLASSES

    def forward(self, input_ids, attention_mask=None):
        """
        实现前向传播过程

        参数:
        input_ids -- 输入的序列标识符，形状为(batch_size, sequence_length)
        attention_mask -- 注意力掩码，用于指示序列中的有效token和填充token，形状为(batch_size, sequence_length)

        返回:
        线性层处理后的输出，用于后续的任务处理
        """
        # 通过BERT模型获取输入的编码表示
        output = self.bert(input_ids=input_ids, attention_mask=attention_mask)

        # 提取BERT模型最后一层隐藏状态
        last_hidden_state = output.last_hidden_state

        # 获取[CLS]标记的输出，用于后续任务的处理
        cls_output = last_hidden_state[:, 0, :]

        # 通过线性层进一步处理[CLS]标记的输出
        return self.linear(cls_output)


class Classifier:
    pass


class ProductClassifier:
    pass